package com.alan6.distributed.ratelimit.handler;

import com.alan6.distributed.ratelimit.annotation.AlaRateLimit;
import com.alan6.distributed.ratelimit.fallback.AlaRateLimitFallback;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import org.springframework.util.Assert;

import javax.annotation.PostConstruct;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;

/**
 * @author ：Alan6
 * @description：TODO
 * @date ：2020/10/13 18:24
 */
@Aspect
@Component
public class RateLimitHandler {

    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimitHandler.class);

    @Autowired
    RedisTemplate redisTemplate;

    private DefaultRedisScript<Long> getRedisScript;

    @PostConstruct
    public void init() {
        getRedisScript = new DefaultRedisScript<>();
        getRedisScript.setResultType(Long.class);
        getRedisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("alaRateLimit.lua")));
        LOGGER.info("AlaRateLimit Script load success");
    }

    @Pointcut("@annotation(com.alan6.distributed.ratelimit.annotation.AlaRateLimit)")
    public void rateLimit() {

    }

    @Around("@annotation(rateLimit)")
    public Object around(ProceedingJoinPoint proceedingJoinPoint, AlaRateLimit rateLimit) throws Throwable {

        Signature signature = proceedingJoinPoint.getSignature();
        if (!(signature instanceof MethodSignature)) {
            throw new IllegalArgumentException("the Annotation @AlaRateLimit must be used on method");
        }
        /**
         * 获取注解参数
         */
        // 限流模块key
        String limitKey = rateLimit.key();

        if (StringUtils.isBlank(limitKey)) {
            // 如果没有设置key，key的默认值=方法包名+方法名
            limitKey = proceedingJoinPoint.getSignature().getDeclaringType().getPackage().getName() + "." + proceedingJoinPoint.getSignature().getName();
        }

        // 限流阈值
        long limitTimes = rateLimit.limit();

        // 限流超时时间
        long expireTime = rateLimit.interval();

        // 限流提示语
        String message = rateLimit.message();

        // 获取降级对象
        Object fallback = rateLimit.fallback().newInstance();


        List<String> keyList = new ArrayList();

        // 设置key值为注解中的值
        keyList.add(limitKey);

        // 调用Lua脚本
        Long result = (Long) redisTemplate.execute(getRedisScript, keyList, expireTime, limitTimes);
        if (result == 0) {
            // 指定了降级
            if (fallback != null) {

                // 检查降级方法是否实现AlaRateLimitFallback接口
                Assert.isTrue(
                        fallback instanceof AlaRateLimitFallback,
                        "Rate limit fallback class must implement the interface AlaRateLimitFallback"
                );

                // 反射调用降级方法
                Method m = fallback.getClass().getMethod("fallbackResult", null);

                // 反射执行方法
                return m.invoke(fallback, null);

            }else {
                return message;
            }
        }

        return proceedingJoinPoint.proceed();
    }
}

